package gpucore

import (
	"fmt"

	"github.com/cogentcore/webgpu/wgpu"
)

// ComputeKernel 表示一个可调度的 GPU 计算程序
type ComputeKernel interface {
	EncodeDispatch(
		encoder *wgpu.CommandEncoder,
		workgroupSize [3]uint32,
		buffers []GPUBuffer,
		uniformBuffers []*wgpu.Buffer,
	) error
}

// compiledKernel 表示一个已编译的 GPU 计算 kernel
type compiledKernel struct {
	device        *wgpu.Device
	pipeline      *wgpu.ComputePipeline
	bindLayout    *wgpu.BindGroupLayout
	workgroupSize [3]uint32
}

func (k *compiledKernel) EncodeDispatch(
	encoder *wgpu.CommandEncoder,
	gridSize [3]uint32, // 重命名：实际为 dispatch 的 workgroup 数量
	buffers []GPUBuffer,
	uniformBuffers []*wgpu.Buffer,
) error {
	// 移除错误的 workgroupSize 覆盖逻辑
	var entries []wgpu.BindGroupEntry

	for i, buf := range buffers {
		entries = append(entries, wgpu.BindGroupEntry{
			Binding: uint32(i),
			Buffer:  buf.Raw(),
			Size:    buf.Size(),
		})
	}

	for i, ub := range uniformBuffers {
		entries = append(entries, wgpu.BindGroupEntry{
			Binding: uint32(len(buffers) + i),
			Buffer:  ub,
			Size:    ub.GetSize(),
		})
	}

	bindGroup, err := k.device.CreateBindGroup(&wgpu.BindGroupDescriptor{
		Layout:  k.bindLayout,
		Entries: entries,
	})
	if err != nil {
		return fmt.Errorf("create bind group: %w", err)
	}
	defer bindGroup.Release()

	pass := encoder.BeginComputePass(nil)
	pass.SetPipeline(k.pipeline)
	pass.SetBindGroup(0, bindGroup, nil)

	// 修复：直接使用 gridSize，不再错误计算
	pass.DispatchWorkgroups(gridSize[0], gridSize[1], gridSize[2])
	pass.End()
	return nil
}
